import torch
import matplotlib.pyplot as plt
from torch import nn

x = torch.linspace(0, 1, 20)
x = x.reshape(-1, 1)
y = torch.zeros(x.shape)
y[10:] = 1.

plt.plot(x, y, 'ro')

model = nn.Sequential(
    nn.Linear(1, 64),
    nn.Tanh(),
    nn.Linear(64, 1)
)
criterion = nn.MSELoss()
sgd = torch.optim.SGD(model.parameters(), lr=0.01)

epochs = 10000
for epoch in range(epochs):
    sgd.zero_grad()
    y_predict = model(x)
    loss = criterion(y_predict, y)
    loss.backward()
    print(loss.item())
    sgd.step()

y_predict = model(x)
plt.plot(x.detach().numpy(), y_predict.detach().numpy(), 'b--')
plt.show()
